// jdbcsh - An interactive shell for JDBC databases
// Copyright (C) 2007 Dan Boitnott
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 2
// of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

package jdbcsh;

import au.com.bytecode.opencsv.CSVReader;
import au.com.bytecode.opencsv.CSVWriter;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.sql.*;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Enumeration;
import java.util.HashMap;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang.time.DateFormatUtils;

/**
 *
 * @author Dan Boitnott
 */
public class Dispatcher {

    public class DispatcherException extends Exception {

        public DispatcherException() {
            super();
        }

        public DispatcherException(Throwable err) {
            super(err);
        }

        public DispatcherException(String err) {
            super(err);
        }
    }
    private Connection conn;
    private String conn_url;
    private String conn_user;
    private String conn_password;

    public void connect(String url) throws DispatcherException {
        disconnect();
        conn_url = url;
        conn_user = null;
        conn_password = null;
        conn = getConnection();
    }

    public void connect(String url, String user, String password) throws DispatcherException {
        disconnect();
        conn_url = url;
        conn_user = user;
        conn_password = password;
        conn = getConnection();
    }

    private Connection getConnection(String url, String user, String password) throws Dispatcher.DispatcherException {
        try {            
            autoLoadDriver(url);
            if (user == null) {
                return DriverManager.getConnection(url);
            }
            return DriverManager.getConnection(url, user, password);
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    private Connection getConnection() throws Dispatcher.DispatcherException {
        return getConnection(conn_url, conn_user, conn_password);
    }

    private void assertConnected() throws DispatcherException {
        if (conn == null) {
            throw new DispatcherException("Not connected");
        }
    }

    public void disconnect() {
        if (conn != null) {
            try {
                conn.close();
            } catch (SQLException ex) {
                throw new RuntimeException(ex);
            }
        }
        conn = null;
    }

    private void autoLoadDriver(String url) {
        // See if a driver is available already
        try {
            DriverManager.getDriver(url);
            return;
        } catch (SQLException ex) {
        }

        // Extract the driver key from the URL
        String[] parts = url.split(":");
        if (parts.length < 2) {
            return;
        }
        String k = parts[1];

        // Translate the driver key
        String drv = null;
        if (k.equals("postgresql")) {
            drv = "org.postgresql.Driver";
        } else if (k.equals("oracle")) {
            drv = "oracle.jdbc.driver.OracleDriver";
        }

        // Return if no driver could be translated
        if (drv == null) {
            return;
        }

        // Attempt to load the driver
        try {
            loadDriver(drv);
        } catch (DispatcherException ex) {
        }
    }

    public void sendSQL(String sql) throws DispatcherException {
        assertConnected();

        try {
            Statement stm = conn.createStatement();
            if (stm.execute(sql)) {
                // A ResultSet was produced
                ResultSet res = stm.getResultSet();

                try {
                    ResultFormatter.printResultSet(System.out, res);
                } catch (Exception ex) {
                    throw new RuntimeException(ex);
                }
            } else {
                // An count was produced
                System.err.println(stm.getUpdateCount() + " rows updated.");                
            }

            printWarnings(stm.getWarnings());
            
            System.out.flush();
            System.err.flush();
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    public void desc(String o, boolean verbose) throws DispatcherException {
        assertConnected();

        try {
            // Split schema if any
            String[] splo = splitObjectSchema(o);
            String obj = splo[1];
            String schema = splo[0];

            // Fix the case for the object & schema
            obj = objectCase(obj);
            schema = objectCase(schema);

            DatabaseMetaData meta = conn.getMetaData();

            // Get the table info
            ResultSet tres = meta.getTables(null, schema, obj, null);
            ResultSet cres = meta.getColumns(null, schema, obj, null);
            ResultSet ires = meta.getIndexInfo(null, schema, obj, false, true);

            if (verbose) {
                ResultFormatter.printResultSet(tres);
                System.out.print("\n");
                ResultFormatter.printResultSet(cres);
                System.out.print("\n");
                System.out.print("\n");
                ResultFormatter.printResultSet(ires);
                return;
            }

            // See if a table was returned
            if (!tres.next()) {
                throw new DispatcherException("object not found");
            }

            // print the table type, name, and schema
            System.out.println(tres.getString("TABLE_TYPE") + " \"" +
                    tres.getString("TABLE_SCHEM") + "." +
                    tres.getString("TABLE_NAME") + "\"\n");

            // print the column details
            ResultFormatter.printTable(new DescTable(cres), false);
            
            // Return here if there are no indexes
            if (!ires.next()) {
                return;
            }

            System.out.println("\nIndexes:");

            String v;
            do {
                v = "    \"" + ires.getString("INDEX_NAME") + "\" ";
                if (!ires.getBoolean("NON_UNIQUE")) {
                    v += "UNIQUE ";
                }
                switch (ires.getShort("TYPE")) {
                case DatabaseMetaData.tableIndexClustered:
                    v += "clustered ";
                    break;
                case DatabaseMetaData.tableIndexHashed:
                    v += "hashed ";
                    break;
                case DatabaseMetaData.tableIndexStatistic:
                    v += "statistic ";
                    break;
                default:
                    v += "other ";
                }
                v += "on (" + ires.getString("COLUMN_NAME") + ")";
                System.out.println(v);
            } while (ires.next());
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    private String[] splitObjectSchema(String o) throws Dispatcher.DispatcherException {
        // Split schema if any
        String obj;
        String schema;
        String[] parts = o.split("\\.");
        if (parts.length == 1) {
            // No schema specified
            obj = o;
            schema = null;
        } else if (parts.length == 2) {
            // A schema was specified
            obj = parts[1];
            schema = parts[0];
        } else {
            // A malformed object was specified
            throw new DispatcherException("Invalid object specifier: " + o);
        }

        String[] ret = {schema, obj};
        return ret;
    }

    public void objectSearch(String schema, String[] types, boolean verbose) throws Dispatcher.DispatcherException {
        try {
            assertConnected();

            DatabaseMetaData meta = conn.getMetaData();
            ResultSet res = meta.getTables(null, objectCase(schema), null, types);

            if (verbose) {
                ResultFormatter.printResultSet(res);
                return;
            }

            FilteredResultSetTable tbl = new FilteredResultSetTable(res);
            tbl.addColumn("TABLE_TYPE", "type");
            tbl.addColumn("TABLE_SCHEM", "schema");
            tbl.addColumn("TABLE_NAME", "name");

            // Pre-load the table for nice column widths
            PreloadTable ptbl = new PreloadTable(tbl);

            ResultFormatter.printTable(ptbl);
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    public void objectSearch(String schema, String type, boolean verbose) throws DispatcherException {
        String[] t = {type};
        objectSearch(schema, t, verbose);
    }

    public void objectSearch(String schema, boolean verbose) throws Dispatcher.DispatcherException {
        String[] t = null;
        objectSearch(schema, t, verbose);
    }

    public String objectCase(String o) throws Dispatcher.DispatcherException {
        if (o == null) {
            return null;
        }

        try {
            assertConnected();
            DatabaseMetaData meta = conn.getMetaData();

            if (isQuotedIdentifier(o)) {
                if (meta.storesUpperCaseQuotedIdentifiers()) {
                    return o.toUpperCase();
                }
                if (meta.storesLowerCaseQuotedIdentifiers()) {
                    return o.toLowerCase();
                }
            } else {
                if (meta.storesUpperCaseIdentifiers()) {
                    return o.toUpperCase();
                }
                if (meta.storesLowerCaseIdentifiers()) {
                    return o.toLowerCase();
                }
            }
            return o;
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    public boolean isQuotedIdentifier(String o) throws Dispatcher.DispatcherException {
        try {
            assertConnected();
            return o.startsWith(conn.getMetaData().getIdentifierQuoteString());
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        }
    }

    private void printWarnings(SQLWarning first) {
        SQLWarning w = first;
        while (w != null) {
            System.err.println("Warning: " + w.getMessage());
            w = w.getNextWarning();
        }
    }

    public void dumpDriverList() {
        for (Enumeration e = DriverManager.getDrivers(); e.hasMoreElements();) {
            System.err.println(e.nextElement());
        }
    }

    public void loadDriver(String drv) throws Dispatcher.DispatcherException {
        System.err.println("Loading JDBC driver: " + drv);
        try {
            Class.forName(drv);
        } catch (ClassNotFoundException ex) {
            throw new DispatcherException("JDBC driver not found.");
        }
        System.err.println("Driver loaded.");
    }
    
    public void dumpCSV(String table, String fpath) throws Dispatcher.DispatcherException {
        try {
            assertConnected();

            // Initialize the writer
            CSVWriter writer = new CSVWriter(new FileWriter(fpath), ',');
            
            // Query the database
            Statement stm = conn.createStatement();
            ResultSet res = stm.executeQuery("select * from " + table);
            
            // Output the CSV
            writer.writeAll(res, true);
                    
            // Finish up
            writer.close();
            res.close();
        } catch (SQLException ex) {
            throw new DispatcherException(ex);
        } catch (IOException ex) {
            throw new DispatcherException(ex);
        }
    }
    
    public void loadCSV(String fpath, String table, String[] columns) throws DispatcherException {
        assertConnected();

        int linenum = 0;
        Connection subconn = null;
        
        try {
            // Initialize the reader
            CSVReader reader = new CSVReader(new FileReader(fpath));
            
            // Read the first line of the file for look-ahead purposes
            String[] line = reader.readNext();
            if (line == null) throw new DispatcherException("file is empty");
            
            String[] splo;
            String schema;
            String obj;
            if (table == null) {
                // No table was specified, create one
                table = "jdbcsh_import_" + DateFormatUtils.format(Calendar.getInstance().getTime(), "yyyyMMddhhmmss");
                schema = null;
                obj = table;
                System.err.println("NOTICE: Creating table " + table);
                String fieldlist = "line_number integer";
                columns = new String[line.length];
                String fname;
                for (int i = 0; i < line.length; i++) {                    
                    fname = "f" + i;
                    fieldlist += ", " + fname + " varchar(255)";
                    columns[i] = fname;
                }
                qexec("create table " + table + " (" + fieldlist + ")");
            } else {
                // Parse the table name
                splo = splitObjectSchema(table);
                schema = splo[0];
                obj = splo[1];
            }
            
            schema = objectCase(schema);
            obj = objectCase(obj);            

            // Get some meta information
            DatabaseMetaData meta = conn.getMetaData();

            // Verify that the table exists            
            ResultSet res = meta.getTables(null, schema, obj, null);
            if (!res.next()) {
                throw new DispatcherException("Table " + table + " does not exist");
            }
            res.close();

            // Determine column quoting
            res = meta.getColumns(null, schema, obj, null);
            HashMap<String, String> quotemap = new HashMap<String, String>();
            while (res.next()) {
                quotemap.put(res.getString("COLUMN_NAME").toLowerCase(), 
                        getQuoting(res.getInt("DATA_TYPE")));
            }            

            // Build the column list                 
            ArrayList<String> clist = new ArrayList<String>();
            String[] quotes = new String[columns.length];
            String col;
            for (int i = 0; i < columns.length; i++) {                
                // Lower the column name's case for matching purposes
                col = columns[i].toLowerCase();                                
                
                // Skip dashed columns
                if (col.equals("-")) continue;
                
                // Verify that the column exists
                if (!quotemap.containsKey(col)) 
                    throw new DispatcherException("Column " + col + " not found in " + table);
                
                clist.add(columns[i]);
                quotes[i] = quotemap.get(col);
            }
            String columnlist = StringUtils.join(clist, ", ");


            // Create a new database connection for transaction isolation
            subconn = getConnection();
            subconn.setAutoCommit(false);

            // Read the file into the table                        
            String sql;
            ArrayList<String> values = new ArrayList<String>();
            do {            
                linenum++;                

                // Verify that the number of fields matches
                if (line.length < columns.length) {
                    throw new DispatcherException("Line " + linenum + " has " +
                            line.length + " fields, expected " + columns.length);
                }
                
                // Build the value list
                values.clear();
                for (int i = 0; i < columns.length; i++) {
                    if (columns[i].equals("-")) continue;
                    values.add(quotes[i] + escapeSQL(line[i].trim()) + quotes[i]);
                }
                String vlist = StringUtils.join(values, ", ");

                // Build the SQL statement
                sql = "insert into " + table + " (" + columnlist + ") values (" + vlist + ")";
                qexec(sql);                
            } while ((line = reader.readNext()) != null);

            // Commit the transaction
            subconn.commit();
        } catch (SQLException ex) {
            if (linenum > 0) {
                throw new DispatcherException("Error on line " + linenum + ": " + ex.getMessage());
            } else {
                throw new DispatcherException(ex);
            }
        } catch (FileNotFoundException ex) {
            throw new DispatcherException("File not found: " + fpath);
        } catch (IOException ex) {
            throw new DispatcherException(ex);
        } finally {
            try {
                if (subconn != null) {
                    subconn.close();
                }
            } catch (SQLException ex) {
            }
        }
    }

    private String getQuoting(int type) {
        switch (type) {
        case Types.BINARY:
        case Types.BLOB:
        case Types.CHAR:
        case Types.CLOB:
        case Types.DATALINK:
        case Types.DATE:
        case Types.JAVA_OBJECT:
        case Types.STRUCT:
        case Types.TIME:
        case Types.TIMESTAMP:
        case Types.VARBINARY:
        case Types.VARCHAR:
            return "'";
        default:
            return "";
        }
    }
    
    private String escapeSQL(String s) {
        return s.replace("'", "''");
    }
    
    private void qexec(String sql) throws SQLException {
        qexec(conn, sql);
    }
    
    private void qexec(Connection sconn, String sql) throws SQLException {
        Statement stm = sconn.createStatement();
        stm.execute(sql);
        stm.close();
    }
}
